import os
from matplotlib import pyplot as plt

from pickle_utils import load_variable


def save_plot(_results, _labels, _filepath):
    fig, ax = plt.subplots()
    for idx, result in enumerate(_results):
        ax.plot(result['repeating_nums'], result['time_consumes'], label=_labels[idx], linewidth=2.0, marker='o')
    ax.set_title('Sort time consumption on datasets with\ndifferent num of repeated elements')
    ax.set_xlabel('num of repeated elements in dataset')
    ax.set_ylabel('consumption (ms)')
    ax.legend()

    # #### need to adjust ####
    ax.set(xlim=(0, 10 ** 6), xticks=range(0, 10 ** 6 + 1, 10 ** 5),
           ylim=(0, 10500), yticks=range(0, 10 * 1000 + 1, 1000))
    # ax.set(xlim=(0, 10 ** 6), xticks=range(0, 10 ** 6 + 1, 10 ** 5),
    #        ylim=(0, 8000), yticks=range(0, 8 * 1000 + 1, 1000))

    plt.savefig(_filepath)


if __name__ == '__main__':
    # result_names = ['quicksort_v1', 'quicksort_v2', 'quicksort_v3']
    result_names = ['quicksort_v3', 'quicksort_v4', 'quicksort_v4_1']
    # result_names = ['quicksort_v2', 'quicksort_v3']
    # labels = ['v1', 'v2', 'v3']
    # labels = ['v2', 'v3', 'v4', 'v4.1']
    labels = ['repeated elements randomized quicksort', '3-way quicksort', 'optimized 3-way quicksort']
    # labels = ['randomize by random numbers', 'randomize by binary flag']
    output_name = 'result.svg'
    # output_name = 'randomize.svg'
    results = [load_variable(os.path.join('results', each, 'time_consumes.pkl')) for each in result_names]
    save_plot(results, labels, f'results/{output_name}')
